package edu.vt.cbil.core;

import edu.vt.cbil.dataset.BioCollection;
import edu.vt.cbil.dataset.BioSample;
import edu.vt.cbil.dataset.ProbeSet;
import edu.vt.cbil.dataset.Segment;
import edu.vt.cbil.util.BioToolbox;
import edu.vt.cbil.util.ConfigObject;
import edu.vt.cbil.util.Histogram;
import org.apache.commons.math3.stat.correlation.PearsonsCorrelation;
import umontreal.iro.lecuyer.probdist.ChiSquareDist;
import umontreal.iro.lecuyer.probdist.ChiSquareNoncentralDist;

import java.util.*;
import java.util.logging.Logger;

/**
 * The type Deletion classifier.
 *
 * @author Justin
 * @version 1.0
 * @created 25 -11月-2014 9:48:26
 */
public class DeletionClassifier {

    /**
     * The Normal rou.
     */
    public double NormalRou;
    private ConfigObject p_config;
    private Logger m_log;
    private BioCollection p_Collection;
    private String curSampleId;
    private BioSample p_sample;


    /**
     * Instantiates a new Deletion classifier.
     *
     * @param pc the pc
     */
    public DeletionClassifier(BioCollection pc, ConfigObject pcfg) {

        m_log = Logger.getLogger("edu.vt.cbil.bacom2");
        this.p_Collection = pc;
        p_config = pcfg;
    }

    /**
     * Sets current sample.
     *
     * @param sampleId the sample id
     */
    public void setCurrentSample(String sampleId) {

        curSampleId = sampleId;
        p_sample = p_Collection.getSample(sampleId);
    }

    /**
     * Detect deletion segments.
     *
     * @param sampleId the sample id
     */
    public void detectDeletionSegment(String sampleId) {

        setCurrentSample(sampleId);

        calculateNoramlABRou();
        calculateY_Lambda();
        boolean existDeletion = bayesianPosterior();

//        double alpha = 0;
        double alpha = calculateAlphaEx(p_sample.fourfoldPeak);
//        if (existDeletion) {
//            alpha = calculateAlpha();
//        } else if (p_sample.fourfoldPeak > 0) {
//            double a = calculateAlphaByFourthPeak(p_sample.fourfoldPeak);
//            if ((a > 0) && (a <= 1)) {
//                alpha = a;
//            }
//        }
        normalizeByPurity(alpha);
    }

    /**
     * Sets collection.
     *
     * @param p_Collection the p _ collection
     */
    public void setCollection(BioCollection p_Collection) {

        this.p_Collection = p_Collection;
    }

    /**
     * Calculate AB-Type Locis's rou in each noraml Segment.
     */
    public void calculateNoramlABRou() {

        List<Double> AlleleA = new ArrayList<Double>();
        List<Double> AlleleB = new ArrayList<Double>();

        for (Short i : p_Collection.ChrIdInterested) {
            for (Segment s : p_sample.getAllSegments(i)) {
                Map<Long, ProbeSet> mps = p_sample.getAllProbeSets(i);

                if (s.fVector.isNeutral) {

                    for (Long id : s.Loci.keySet()) {
                        ProbeSet p = mps.get(id);
                        if ((p.Is_SNPProbe) && (p.data.Normal.isGenotypeAB) && p.data.isNotLOH) {
                            if (p_config.params.enableCorrectCrossTalk) {
                                AlleleA.add(p.data.alleleA_CorrectCrossTalk);
                                AlleleB.add(p.data.alleleB_CorrectCrossTalk);
                            } else {
                                AlleleA.add(p.data.alleleA);
                                AlleleB.add(p.data.alleleB);
                            }
                        }
                    }
                }
            }
        }
        //convert from Double[] to double[]
        int SNPlen = AlleleA.size();
        if (SNPlen > 3) {
            double[] a = new double[SNPlen];
            double[] b = new double[SNPlen];
            for (int j = 0; j < SNPlen; j++) {
                a[j] = AlleleA.get(j);
                b[j] = AlleleB.get(j);
            }

            PearsonsCorrelation pCorr = new PearsonsCorrelation();
            NormalRou = pCorr.correlation(a, b);

        } else {
            NormalRou = 1.0;
        }
        m_log.info(String.format("Rou of copy-neutral Loci= %f; SNP count = %d", NormalRou, SNPlen));

    }

    /**
     * Calculate Y and lambda.
     */
    public void calculateY_Lambda() {

        List<Double> AlleleA = new ArrayList<Double>();
        List<Double> AlleleB = new ArrayList<Double>();
        List<Double> CopyNumber = new ArrayList<Double>();

        for (Short i : p_Collection.getChrIdArray()) {
            for (Segment s : p_sample.getAllSegments(i)) {
                Map<Long, ProbeSet> mps = p_sample.getAllProbeSets(i);

                if (s.fVector.isDeletion) {
                    AlleleA.clear();
                    AlleleB.clear();
                    CopyNumber.clear();

                    for (Long id : s.Loci.keySet()) {
                        ProbeSet p = mps.get(id);
                        if ((p.Is_SNPProbe) && (p.data.Normal.isGenotypeAB)) {
                            AlleleA.add(2 * p.data.alleleA);
                            AlleleB.add(2 * p.data.alleleB);

                            CopyNumber.add(2 * p.data.CopyNumberHalf); //same
                            //CopyNumber.add(2 * (p.data.alleleA + p.data.alleleB));
                        }
                    }

                    int len = AlleleA.size();
                    s.fVector.SNPABLength = len;
                    if (s.fVector.isIgnored())
                        continue;

                    double[] tempSNPValue = BioToolbox.robustMean(CopyNumber, 64);
                    //s.fVector.robustMeanSNP_AB_CN = tempSNPValue[0];
                    s.fVector.robustStdSNP_AB_CN = tempSNPValue[1];
                    s.fVector.robustMeanSNP_AB_CN = 2 * s.robustMeanHalfCopyNumber;

                    //sigma(A+B)
                    double sigmaSumAB = Math.pow(s.fVector.robustStdSNP_AB_CN, 2);
                    //sigma(A-B)
                    double sigmaDiffAB = sigmaSumAB * (1 - NormalRou) / (1 + NormalRou);
                    double score;

                    Iterator<Double> itA = AlleleA.iterator();
                    Iterator<Double> itB = AlleleB.iterator();
                    score = 0;

                    for (int j = 0; j < len; j++) {
                        score += Math.pow(itA.next() - itB.next(), 2);
                    }

                    s.fVector.YValue = score / sigmaDiffAB;
                    s.fVector.Lambda = len * Math.pow(2 - s.fVector.robustMeanSNP_AB_CN, 2) / sigmaDiffAB;
                }

            }
            //System.out.println(i);
        }

    }

    /**
     * Bayesian posterior.
     */
    public boolean bayesianPosterior() {

        double[] SumAB_CN_Seg = new double[2];// the sum of CN of all AB-Type Locis , first is hemi
        int[] CountAB_Seg = new int[2];
        double[] segMean = new double[2];

        for (Short i : p_Collection.getChrIdArray()) {
            for (Segment s : p_sample.getAllSegments(i)) {

                if (s.fVector.isIgnored())
                    continue;

                if (s.fVector.isDeletion) {

                    int segSize = s.fVector.SNPABLength;

                    ChiSquareDist homoDist = new ChiSquareDist(segSize - 1);

                    ChiSquareNoncentralDist hemiDist = new ChiSquareNoncentralDist(segSize - 1, s.fVector.Lambda);

                    double homop = homoDist.density(s.fVector.YValue);
                    double hemip = hemiDist.density(s.fVector.YValue);

//
//                    if (homop < 1e-12) homop = 0;
//                    if (hemip < 1e-12) hemip = 0;

                    //if ((homop == 0) && (hemip == 0)) {
//                    if (s.Loci.size() < 64) {
//                        s.fVector.isHomoDeletion = false;
//                        s.fVector.isHemiDeletion = false;
//
//                        m_log.info(String.format("Chr_id = %2d \t avg = %.6f \t Ignored",
//                                s.Chr_id, s.fVector.robustMeanSNP_AB_CN));
//
//                        continue;
//                    }
                    if (homop == 0 && hemip == 0) continue;
                    s.fVector.PValue_Homo = homop;
                    s.fVector.PValue_Hemi = hemip;


                    if (s.fVector.YValue <= s.fVector.SNPABLength) {
                        s.fVector.isHomoDeletion = true;
                        s.fVector.isHemiDeletion = false;
                    } else if (s.fVector.YValue >= s.fVector.SNPABLength + s.fVector.Lambda) {
                        s.fVector.isHomoDeletion = false;
                        s.fVector.isHemiDeletion = true;
                    } else if (s.fVector.PValue_Homo > s.fVector.PValue_Hemi) {
                        s.fVector.isHomoDeletion = true;
                        s.fVector.isHemiDeletion = false;
                    } else {
                        s.fVector.isHomoDeletion = false;
                        s.fVector.isHemiDeletion = true;
                    }

                    if (s.fVector.isHemiDeletion) {
                        SumAB_CN_Seg[0] += s.fVector.robustMeanSNP_AB_CN * s.fVector.SNPABLength;
                        CountAB_Seg[0] += s.fVector.SNPABLength;
                    } else {
                        SumAB_CN_Seg[1] += s.fVector.robustMeanSNP_AB_CN * s.fVector.SNPABLength;
                        CountAB_Seg[1] += s.fVector.SNPABLength;
                    }
//                    Map<Long, ProbeSet> mps = p_Collection.sample.get(i).probes;
//                    for (Long id : s.Loci.keySet()){
//                        ProbeSet ps = mps.get(id);
//                        if (ps.Is_SNPProbe && ps.data.Normal.isGenotypeAB){
//                            if (s.fVector.isHemiDeletion) {
//                                segSum[0] += ps.data.CopyNumberHalf;
//                                segCount[0]++;
//                            } else {
//                                segSum[1] += ps.data.CopyNumberHalf;
//                                segCount[1]++;
//                            }
//                        }
//                    }
                }
            }
        }

        if ((CountAB_Seg[0] == 0) && (CountAB_Seg[1] == 0)) {
            m_log.severe("don't exist any deletions!");
            return false;

        } else {
            if (CountAB_Seg[0] > 0) {
                segMean[0] = SumAB_CN_Seg[0] / CountAB_Seg[0];
            } else {
                segMean[0] = 2;
            }
            if (CountAB_Seg[1] > 0) {
                segMean[1] = SumAB_CN_Seg[1] / CountAB_Seg[1];
            } else {
                segMean[1] = 0;
            }
            m_log.info(String.format("Hemi: Loci Count = % 5d \t Mean = %.5f", CountAB_Seg[0], segMean[0]));
            m_log.info(String.format("Homo: Loci Count = % 5d \t Mean = %.5f", CountAB_Seg[1], segMean[1]));

            double thresh;
            if ((segMean[0] > segMean[1]) && (CountAB_Seg[1] > 0)) {
                thresh = (segMean[0] + segMean[1]) / 2.0;
            } else {
                thresh = 0.95 * segMean[0];
            }
            m_log.info(String.format("Homo thresh = %.5f", thresh));

            for (Short i : p_Collection.getChrIdArray()) {
                for (Segment s : p_sample.getAllSegments(i)) {

                    if (s.fVector.isIgnored()) continue;

                    if (s.fVector.isDeletion) {
                        if (s.fVector.robustMeanSNP_AB_CN < thresh) {
                            s.fVector.isHomoDeletion = true;
                            s.fVector.isHemiDeletion = false;
                        } else {
                            s.fVector.isHomoDeletion = false;
                            s.fVector.isHemiDeletion = true;
                        }
                    }
                }
            }
            return true;
        }
    }

    private double calculateAlphaByMean() {

        double sum = 0;
        int count = 0;
        for (Short i : p_Collection.getChrIdArray()) {
            for (Segment s : p_sample.getAllSegments(i)) {

                if (s.fVector.isIgnored()) continue;

                if (s.fVector.isHomoDeletion) {
                    s.fVector.alpha = s.fVector.robustMeanSNP_AB_CN / 2.0;
                } else if (s.fVector.isHemiDeletion) {
                    s.fVector.alpha = s.fVector.robustMeanSNP_AB_CN - 1;
                }

                sum += s.fVector.alpha * s.Loci.size();
                count += s.Loci.size();
//                sum += s.fVector.alpha*s.fVector.SNPABLength;
//                count += s.fVector.SNPABLength;

                if (s.fVector.isHomoDeletion) {
                    m_log.info(String.format("Chr_id = %2d \t count = % 6d \t avg = %.6f \t *homo = %.6e \t  hemi = %.6e \t alpha = %.6f",
                            s.Chr_id, s.fVector.SNPABLength, s.fVector.robustMeanSNP_AB_CN, s.fVector.PValue_Homo, s.fVector.PValue_Hemi, s.fVector.alpha));
                } else {
                    m_log.info(String.format("Chr_id = %2d \t count = % 6d \t avg = %.6f \t  homo = %.6e \t *hemi = %.6e \t alpha = %.6f",
                            s.Chr_id, s.fVector.SNPABLength, s.fVector.robustMeanSNP_AB_CN, s.fVector.PValue_Homo, s.fVector.PValue_Hemi, s.fVector.alpha));
                }
            }
        }

        if (count > 0) {
            return sum / count;
        } else {
            return 0;
        }

    }

    private double calculateAlphaByFourthPeak(double peak) {


        double upLimit = (1 + p_config.params.widthFourthPart) * peak;
        double downLimit = (1 - p_config.params.widthFourthPart) * peak;

        List<Double> ts = new ArrayList<Double>();

        for (Short i : p_Collection.getChrIdArray()) {
            for (Segment s : p_sample.getAllSegments(i)) {

                if ((s.robustMeanHalfCopyNumber < upLimit) && (s.robustMeanHalfCopyNumber > downLimit)) {
                    s.fVector.alpha = 2 - s.robustMeanHalfCopyNumber;
                    ts.add(s.fVector.alpha);
                }
            }
        }

        if (ts.size() > 0) {
            Collections.sort(ts);
            int pos = (int) (ts.size() * p_config.params.ALPHA_POSITION);
            return ts.get(pos);
        }

        return 0;
    }

    private double calculateAlphaEx(double peak) {

        double upLimit = 1.1 * peak;
        double downLimit = 0.9 * peak;
        Histogram h = new Histogram(16, 0, 1);
        List<Double> ts = new ArrayList<Double>();

        for (Short i : p_Collection.getChrIdArray()) {
            for (Segment s : p_sample.getAllSegments(i)) {

                if ((peak > 0) && (s.robustMeanHalfCopyNumber < upLimit) && (s.robustMeanHalfCopyNumber > downLimit)) {
                    s.fVector.alpha = 2 - s.robustMeanHalfCopyNumber;
                } else {
                    if (s.fVector.isHomoDeletion) {
                        s.fVector.alpha = s.fVector.robustMeanSNP_AB_CN / 2.0;
                    } else if (s.fVector.isHemiDeletion) {
                        s.fVector.alpha = s.fVector.robustMeanSNP_AB_CN - 1;
                    }
                }

                if (s.fVector.alpha > 0.05 && s.fVector.alpha < 0.95) {
                    ts.add(s.fVector.alpha);
                    h.addDataPoint(s.fVector.alpha, s.Loci.size());
                }
            }
        }
        m_log.info(h.toString("alpha "));

        if (ts.size() > 0) {
            Collections.sort(ts);
            int pos = (int) (ts.size() * p_config.params.ALPHA_POSITION);
            return ts.get(pos);
        }

        return 0;
    }

    private double calculateAlpha() {

        Histogram h = new Histogram(16, 0, 1);
        List<Double> ts = new ArrayList<Double>();

        for (Short i : p_Collection.getChrIdArray()) {
            for (Segment s : p_sample.getAllSegments(i)) {

                if (s.fVector.isIgnored()) continue;

                if (s.fVector.isHomoDeletion) {
                    s.fVector.alpha = s.fVector.robustMeanSNP_AB_CN / 2.0;
                } else if (s.fVector.isHemiDeletion) {
                    s.fVector.alpha = s.fVector.robustMeanSNP_AB_CN - 1;
                }

                h.addDataPoint(s.fVector.alpha);
                ts.add(s.fVector.alpha);


                if (s.fVector.isHomoDeletion) {
                    m_log.info(String.format("Chr_id = %2d \t count = % 6d \t avg = %.6f \t *homo = %.6e \t  hemi = %.6e \t alpha = %.6f",
                            s.Chr_id, s.fVector.SNPABLength, s.fVector.robustMeanSNP_AB_CN, s.fVector.PValue_Homo, s.fVector.PValue_Hemi, s.fVector.alpha));
                } else {
                    m_log.info(String.format("Chr_id = %2d \t count = % 6d \t avg = %.6f \t  homo = %.6e \t *hemi = %.6e \t alpha = %.6f",
                            s.Chr_id, s.fVector.SNPABLength, s.fVector.robustMeanSNP_AB_CN, s.fVector.PValue_Homo, s.fVector.PValue_Hemi, s.fVector.alpha));
                }
            }
        }


        m_log.info(h.toString("s.fv.alpha "));

        if (ts.size() > 0) {
            Collections.sort(ts);
            int pos = (int) (ts.size() * p_config.params.ALPHA_POSITION);
            return ts.get(pos);
        } else {
            m_log.warning("Hemi and Homo also are null!");
            return 0;
        }
    }

    private void normalizeByPurity(double alpha) {

        double purity = 1 - alpha;

        for (Short i : p_Collection.getChrIdArray()) {
            for (Segment s : p_sample.getAllSegments(i)) {
                s.purityCorrCopyNumber = 2 * (s.robustMeanHalfCopyNumber - alpha) / purity;
                if (s.purityCorrCopyNumber < 1.7) {
                    s.fVector.isDeletion = true;
                    s.fVector.isNeutral = false;
                } else if (s.purityCorrCopyNumber > 2.3) {
                    s.fVector.isDeletion = false;
                    s.fVector.isNeutral = false;
                }
            }
        }

        p_Collection.results.setPurity(curSampleId, purity);

        m_log.info(String.format("alpha = %.5f", 1 - purity));
        m_log.info(String.format("1 - alpha = %.5f", purity));
    }

//    private void normalizeByPurity(double alpha) {
//        double purity = 1 - alpha;
//        for (Short i : p_Collection.getChrIdArray()) {
//            for (Segment s : p_sample.getAllSegments(i)) {
//                s.purityCorrCopyNumber = 2 * (s.robustMeanHalfCopyNumber - alpha) / purity;
//            }
//        }
//
//        double sumHemi = 0;
//        int countHemi = 0;
//        double sumHomo = 0;
//        int countHomo = 0;
//        for (Short i : p_Collection.getChrIdArray()) {
//            for (Segment s : p_sample.getAllSegments(i)) {
//                int len = s.Loci.size();
//                if (s.fVector.isHemiDeletion) {
//                    sumHemi += s.purityCorrCopyNumber * len;
//                    countHemi += len;
//                } else if (s.fVector.isHomoDeletion) {
//                    sumHomo += s.purityCorrCopyNumber * len;
//                    countHomo += len;
//                }
//            }
//        }
//        double avgHemi = (countHemi>0)? sumHemi / countHemi:0;
//        double avgHomo = (countHomo>0)? sumHomo / countHomo:0;
//
//        m_log.info(String.format("Hemi: probeCount = % 5d \t Mean = %.5f", countHemi, avgHemi));
//        m_log.info(String.format("Homo: probeCount = % 5d \t Mean = %.5f", countHomo, avgHomo));
//
//        if ((avgHemi - avgHomo < 0.5) && (countHemi + countHomo > 0)) {
//            double k = (2 - (sumHemi+sumHomo) / (countHemi+countHomo));
//            m_log.info(String.format("k = %.5f", k));
//            purity = purity*k;
//
//            for (Short i : p_Collection.getChrIdArray()) {
//                for (Segment s : p_sample.getAllSegments(i)) {
//                    s.purityCorrCopyNumber = (s.purityCorrCopyNumber - 2) / k + 2;
//                }
//            }
//        }
//
//        p_Collection.results.setPurity(curSampleId, purity);
//
//        m_log.info(String.format("alpha = %.5f", 1 - purity));
//        m_log.info(String.format("1 - alpha = %.5f", purity));
//    }


}

